from langchain_core.embeddings import Embeddings
import requests
import json
import os

class AliyunDashScopeEmbeddings(Embeddings):
    def __init__(self, model: str = "text-embedding-v4"):
        self.model = model
        self.dashscope_api_key = os.getenv("API_KEY")
        if not self.dashscope_api_key:
            raise ValueError("Did not find DashScope API key, please set DASHSCOPE_API_KEY environment variable.")

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        url = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding"
        headers = {
            "Authorization": f"Bearer {self.dashscope_api_key}",
            "Content-Type": "application/json"
        }

        data = {
            "model": self.model,
            "input": {
                "texts": texts
            }
        }

        response = requests.post(url, headers=headers, data=json.dumps(data))
        if response.status_code == 200:
            result = response.json()
            embeddings = [item["embedding"] for item in result["output"]["embeddings"]]
            return embeddings
        else:
            raise Exception(f"Failed to get embeddings: {response.text}")

    def embed_query(self, text: str) -> list[float]:
        return self.embed_documents([text])[0]